#pragma once
#include "sm2.h"
#include "DataVerifyOps.h"
#include <windows.h>


Crypto_SM2Key sm2_key = {
    .privateKey = {0x39, 0x45, 0x20, 0x8F, 0x7B, 0x21, 0x44, 0xB1, 0x3F, 0x36, 0xE3, 0x8A,
        0xC6, 0xD3, 0x9F, 0x95, 0x88, 0x93, 0x93, 0x69, 0x28, 0x60, 0xB5, 0x1A,
        0x42, 0xFB, 0x81, 0xEF, 0x4D, 0xF7, 0xC5, 0xB8
        },
    .publicKeyX = {0x09, 0xF9, 0xDF, 0x31, 0x1E, 0x54, 0x21, 0xA1, 0x50, 0xDD, 0x7D, 0x16,
        0x1E, 0x4B, 0xC5, 0xC6, 0x72, 0x17, 0x9F, 0xAD, 0x18, 0x33, 0xFC, 0x07,
        0x6B, 0xB0, 0x8F, 0xF3, 0x56, 0xF3, 0x50, 0x20
         },
    .publicKeyY = {0xCC, 0xEA, 0x49, 0x0C, 0xE2, 0x67, 0x75, 0xA5, 0x2D, 0xC6, 0xEA, 0x71,
        0x8C, 0xC1, 0xAA, 0x60, 0x0A, 0xED, 0x05, 0xFB, 0xF3, 0x5E, 0x08, 0x4A,
        0x66, 0x32, 0xF6, 0x07, 0x2D, 0xA9, 0xAD, 0x13
        }
};


/**
 * Calculate SM2 hash message digest
 *
 * @param zValue concatenate constants and hash them yields zValue
 * @param message data awaiting signature
 * @param msgLen the length of the message
 * @param digest the result of concatenating messages with z-values and hashing them
 */
bool SM2_HashMessage(const uint8_t* zValue, const uint8_t* message, size_t msgLen, uint8_t* digest)
{
    size_t zValueLen = 32; // zValue 的长度固定为 32 字节
    size_t totalLen = zValueLen + msgLen;


    // 创建一个临时缓冲区存放 zValue 和消息的拼接
    uint8_t* data = (uint8_t*)malloc(totalLen);
    if (data == NULL || zValue == NULL)
    {
        printf("memery allocate failed!");
        return false; // 内存分配失败
    }

    // 拼接 zValue 和消息
    memcpy(data, zValue, zValueLen);
    memcpy(data + zValueLen, message, msgLen);
    // 打印拼接后的消息值为十六进制
    //printf("拼接的消息值: \n");
    //for (size_t i = 0; i < totalLen; ++i) {
    //    printf("%02X ", data[i]);  // 每个字节输出 2 个十六进制字符
    //}
    //printf("\n");
    // 计算 SM3 哈希值
    sm3(data, totalLen, digest);

    // 释放临时缓冲区
    free(data);

    return true;
}

/**
 * Calculate user identification Z value
 *
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param a elliptical curve parameter a
 * @param b elliptical curve parameter b
 * @param Gx the x-coordinate of the base point G
 * @param Gy the y-coordinate of the base point G
 * @param Px the x-coordinate of the public key
 * @param Py the y-coordinate of the public key
 * @param Z output Z-value (hash result)
 *
 */
bool SM2_compute_Z(
    uint8_t* ID,
    uint16_t ID_len,
    uint8_t* a,
    uint8_t* b,
    uint32_t* Gx,
    uint32_t* Gy,
    uint8_t* Px,
    uint8_t* Py,
    uint8_t* Z
)
{
    sm2_uint256_t _32b_zValue;
    uint8_t _8b_zValue[32];
    size_t buffer_size = 2 + ID_len + 32 * 6;
    uint8_t* buffer = (uint8_t*)malloc(buffer_size);
    if (buffer == NULL) {
        printf("内存分配失败！\n");
        return -1;  // 分配失败
    }
    uint16_t ENTL = ID_len * 8; // ID 长度（比特表示）
    uint32_t pos = 0;
    // Gx,Gy要转换为8位类型
    uint8_t _8b_Gx[32], _8b_Gy[32];
    convert_to_uint8(Gx, _8b_Gx);
    convert_to_uint8(Gy, _8b_Gy);


    // 拼接 ENTL（2 字节，大端序）
    buffer[pos++] = (ENTL >> 8) & 0xFF;
    buffer[pos++] = ENTL & 0xFF;

    // 拼接 ID
    memcpy(buffer + pos, ID, ID_len);
    pos += ID_len;

    // 拼接 a, b, Gx, Gy, Px, Py
    memcpy(buffer + pos, a, 32);
    pos += 32;
    memcpy(buffer + pos, b, 32);
    pos += 32;
    memcpy(buffer + pos, _8b_Gx, 32);
    pos += 32;
    memcpy(buffer + pos, _8b_Gy, 32);
    pos += 32;
    memcpy(buffer + pos, Px, 32);
    pos += 32;
    memcpy(buffer + pos, Py, 32);
    pos += 32;
    sm3(buffer, pos, Z);
    return true; // 假设 sm3_final 总是成功
}


/**
 * Signature function
 *
 * @param key secret key
 * @param message data awaiting signature
 * @param msgLen the length of the message
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param signature signature value for message
 */
bool Crypto_SM2Sign(
    const Crypto_SM2Key* key,
    uint8_t* message,
    uint32_t msgLen,
    uint8_t* ID,
    uint16_t ID_len,
    Crypto_SM2Signature* signature)
{
    if (!key || !message || !signature)
    {
        return false; /* 参数检查 */
    }

    uint8_t zValue[32] = {0};   // Z 值

    uint8_t digest[32] = { 0 };   // 消息摘要
    uint8_t r[32] = { 0 };
    uint8_t s[32] = { 0 }; // 签名结果
    sm2_uint256_t random_k = {
    0x59276E27, 0xD506861A,
    0x16680F3A, 0xD9C02DCC,
    0xEF3CC1FA, 0x3CDBE4CE,
    0x6D54B80D, 0xEAC1BC21
    };

    // 计算 Z 值
    if (!SM2_compute_Z(ID, ID_len, param_A, param_B, param_G.x,
        param_G.y, sm2_key.publicKeyX, sm2_key.publicKeyY, zValue))
    {
        return false;
    }
    //printf("zValue:\n");
    //print_uint8_array(zValue);

    // 拼接 Z 和消息计算摘要
    if (!SM2_HashMessage(zValue, message, msgLen, digest))
    {
        return false;
    }

    printf("Message digest:\n");
    print_uint8_array(digest);
    printf("\n");

    sm2_uint256_t _32b_digest = { 0 };
    sm2_uint256_t _32b_r = { 0 };
    sm2_uint256_t _32b_x = { 0 };
    sm2_uint256_t mg_r = { 0 };

    sm2_point_t mg_G;
    // 随机数 k 和签名计算
    while (1)
    {
        //generate_random(k, 32); /* 生成随机数 */

        // 1. 计算椭圆曲线点 (x1, y1) = [k]G
        sm2_point_t point;
        point_modp2mgmy(&param_G, Prime, &mg_G);
        point_multiply(&mg_G, random_k, Prime, &point);
        mgmy_modp2point(&point, Prime, &point);

        // 2. 计算 r = (e + x1) mod n

        convert_to_uint32(digest, _32b_digest);
        //convert_to_uint32(param_N, _32b_param_N);
        mod_add(_32b_digest, point.x, param_N, _32b_r);
        //printf("r的值:\n");
        //print_uint32_array(_32b_r);

        convert_to_uint8(_32b_r, r);

        // 检查 r 和 (r + k) 是否有效
        sm2_uint256_t tmp_k = {0};
        mod_add(_32b_r, random_k, param_N, tmp_k);
        if (!is_zero(_32b_r) && !is_zero(tmp_k))
        {
            break;
        }
    }

    // 计算 s
    sm2_uint256_t _32b_privateKey = {0};
    sm2_uint256_t mg_privateKey = {0};
    sm2_uint256_t mg_s = { 0 };
    sm2_uint256_t tmp1 = { 0 };
    sm2_uint256_t tmp2 = { 0 };
    convert_to_uint32(key->privateKey, _32b_privateKey);
    n_num2mgmy(_32b_privateKey, param_N, mg_privateKey);
    n_num2mgmy(_32b_r, param_N, mg_r);
    // 1. tmp1 = r * d
    n_mod_mul(mg_r, mg_privateKey, param_N, tmp1);
    sm2_uint256_t _32b_k, mg_k;
    //convert_to_uint32(k, _32b_k);
    n_num2mgmy(random_k, param_N, mg_k);
    // 2. tmp2 = k - tmp1
    mod_sub(mg_k, tmp1, param_N, tmp2);
    // 3. tmp1 = (1 + d)^-1
    sm2_uint256_t mg_privateKey_add_1, mg_one;
    sm2_uint256_t one = { 0,0,0,0,0,0,0,1 };
    n_num2mgmy(one, param_N, mg_one);
    mod_add(mg_privateKey, mg_one, param_N, mg_privateKey_add_1);
    n_mod_inverse(mg_privateKey_add_1, param_N, tmp1);
    // 4. mg_s = tmp1 * tmp1  
    n_mod_mul(tmp1, tmp2, param_N, mg_s);
    n_mgmy2num(mg_s, param_N, mg_s);
    convert_to_uint8(mg_s, s);
    //检查 s 是否有效
    if (is_zero(mg_s))
    {
        return false;
    }


    memcpy(signature->r, r, 32);
    memcpy(signature->s, s, 32);

    return true;
}

/**
 * Verification function
 *
 * @param key secret key
 * @param message data awaiting signature
 * @param msgLen the length of the message
 * @param ID user ID
 * @param ID_len user identification length (in bytes)
 * @param signature signature value for message
 */
bool Crypto_SM2Verify(
    const Crypto_SM2Key* key,
    uint8_t* message,
    uint32_t msgLen,
    uint8_t* ID,
    uint16_t ID_len,
    const Crypto_SM2Signature* signature)
{

    if (!key || !message || !signature)
    {
        return false;
    }

    uint8_t zValue[THIRTY_TRHREE];
    uint8_t digest[THIRTY_TRHREE];
    uint8_t t[THIRTY_TRHREE], u1[THIRTY_TRHREE], u2[THIRTY_TRHREE];
    sm2_point_t point1, point2, point;
    sm2_uint256_t _32b_param_N = { 0 };
    sm2_uint256_t _32b_t = { 0 };
    sm2_uint256_t _32b_r = { 0 };
    sm2_uint256_t _32b_s = { 0 };
    convert_to_uint32(signature->r, _32b_r);
    convert_to_uint32(signature->s, _32b_s);

    // 判断r和s是否在[1, n-1]内
    if (!is_in_N(_32b_r) || !is_in_N(_32b_s))
        return false;

    // 计算 Z 值
    if (!SM2_compute_Z(ID, ID_len, param_A, param_B, param_G.x,
        param_G.y, sm2_key.publicKeyX, sm2_key.publicKeyY, zValue))
    {
        return false;
    }

    // 拼接 Z 和消息计算摘要
    if (!SM2_HashMessage(zValue, message, msgLen, digest))
    {
        return false;
    }

    // 验签步骤
    // t = r + s mod n
    mod_add(_32b_r, _32b_s, param_N, _32b_t);
    if (is_zero(_32b_t))
    {
        return false;
    }
    sm2_uint256_t _32b_pub_x, _32b_pub_y, _32b_digest, _32b_r_check;
    sm2_uint256_t mg_t, mg_r, mg_s, mg_pub_x, mg_pub_y;
    sm2_point_t mg_G, mg_key_point;
    //convert_to_uint32(Prime, P);
    point_modp2mgmy(&param_G, Prime, &mg_G);
    num2mgmy(_32b_r, Prime, mg_r);
    num2mgmy(_32b_s, Prime, mg_s);
    num2mgmy(_32b_t, Prime, mg_t);
    //(x1,x2) = sG + tPa
    point_multiply(&mg_G, _32b_s, Prime, &point1);
    convert_to_uint32(key->publicKeyX, _32b_pub_x);
    convert_to_uint32(key->publicKeyY, _32b_pub_y);
    //num2mgmy(_32b_pub_x, P, mg_pub_x);
    //num2mgmy(_32b_pub_y, P, mg_pub_y);
    sm2_point_t key_point;
    memcpy(key_point.x, _32b_pub_x, sizeof(_32b_pub_x));  // 复制 x 坐标
    memcpy(key_point.y, _32b_pub_y, sizeof(_32b_pub_y));  // 复制 y 坐标

    point_modp2mgmy(&key_point, Prime, &mg_key_point);
    point_multiply(&mg_key_point, _32b_t, Prime, &point2);
    point_add(&point1, &point2, Prime, &point);
    sm2_point_t normal_point;
    mgmy_modp2point(&point, Prime, &normal_point);
    // R = e+x1 mod n
    convert_to_uint32(digest, _32b_digest);
    uint8_t r_check[THIRTY_TRHREE];
    mod_add(_32b_digest, normal_point.x, param_N, _32b_r_check);

    if (compare_value(_32b_r, _32b_r_check) == 0)
        return true;
    else {
        return false;
    }
}


/**
* Print signature
*/
void printSignature(const Crypto_SM2Signature* sig) {
    // 打印 r 值
    printf("r: ");
    for (size_t i = 0; i < 32; ++i) {
        printf("%02x", sig->r[i]);  // 按字节打印，2 位十六进制表示
    }
    printf("\n");

    // 打印 s 值
    printf("s: ");
    for (size_t i = 0; i < 32; ++i) {
        printf("%02x", sig->s[i]);  // 按字节打印，2 位十六进制表示
    }
    printf("\n");
}





/**
* Receive user input
*/
void receive_data_from_user(uint8_t** message, char** userID, uint8_t* private_key, uint8_t* public_key, size_t* message_len, size_t* userID_len) {
    // 接收 message（16进制字节数组）
    printf("Enter your message: ");
    *message = get_message(message_len);
    printf("Received message (length: %zu bytes)\n", *message_len);

    // 接收 userID（字符串）
    printf("Enter your userID: ");
    getline_(userID, userID_len, stdin);  // 动态分配内存并接收输入
    (*userID)[strcspn(*userID, "\n")] = '\0';  // 去掉换行符
    printf("Received message (length: %zu bytes)\n", *userID_len / 8);

    // 接收私钥（16进制32个字节数组）
    printf("Enter your 32-byte private key in hexadecimal format (64 characters): ");
    for (int i = 0; i < THIRTY_TRHREE; i++) {
        unsigned int byte;
        if (scanf_s("%2x", &byte) != 1) {
            printf("Invalid input for private key.\n");
            return;
        }
        private_key[i] = (uint8_t)byte;
    }

    // 接收公钥（16进制32个字节数组）
    printf("Enter your 32-byte public key in hexadecimal format (128 characters): ");
    for (int i = 0; i < 2 * THIRTY_TRHREE; i++) {
        unsigned int byte;
        if (scanf_s("%2x", &byte) != 1) {
            printf("Invalid input for public key.\n");
            return;
        }
        public_key[i] = (uint8_t)byte;
    }
}
int main()
{
    //sm2_uint256_t randnum;
    //if (random_generate(randnum) == -1)
    //{
    //    printf("rand gen fail\n");
    //}
    //print_uint32_array(randnum);
    //return 0;
     uint8_t* message = NULL;
     char* userID = NULL;
     uint8_t private_key[THIRTY_TRHREE] = { 0 };  // 初始化为 0
     uint8_t public_key[2 * THIRTY_TRHREE] = { 0 };   // 初始化为 0
     size_t message_len = 0;
     size_t userID_len = 0;
     // Print the header banner
     printf("**************************************** SM2 Signature & Verification Platform ****************************************\n");
     printf("****************************************    Welcome to SM2 Platform    ****************************************\n\n");

     while (true) {
         int mode = 1;
         // Display options
         printf("Please select the testing mode or manually input values (message, user ID, private key, public key):\n");
         printf("--------------------------------------------------------------------------------------------------------\n");
         printf("1. Use National Standard Example 1 for testing\n");
         printf("2. Manually input values (message, user ID, private key, public key)\n");
         printf("--------------------------------------------------------------------------------------------------------\n");

         // Get user input
         printf("Enter your choice (1 or 2): ");
         scanf_s("%d", &mode);
         if (mode == 2) {
             printf("\nYou selected: Manually input values (message, user ID, private key, public key)\n");

             // 接收消息
             receive_data_from_user(&message, &userID, private_key, public_key, &message_len, &userID_len);
             //receive_data_from_user(message, userID, private_key, public_key, &message_len, &userID_len);

             // memcpy(r->x, x3, sizeof(x3));  // 将 x3 的值赋给 R->x
             memcpy(sm2_key.publicKeyX, public_key, sizeof(private_key));
             memcpy(sm2_key.publicKeyY, public_key + 32, sizeof(private_key));
             memcpy(sm2_key.privateKey, private_key, sizeof(private_key));
         }
         else if (mode != 1 && mode != 2) {
             printf("\nInvalid input, exiting....\n");
             break;
         }
         else {
             printf("\nYou selected: Use National Standard Example 1 for testing\n");
             message = (uint8_t*)malloc(14 * sizeof(uint8_t));
             if (message != NULL) {
                 message[0] = 0x6D;
                 message[1] = 0x65;
                 message[2] = 0x73;
                 message[3] = 0x73;
                 message[4] = 0x61;
                 message[5] = 0x67;
                 message[6] = 0x65;
                 message[7] = 0x20;
                 message[8] = 0x64;
                 message[9] = 0x69;
                 message[10] = 0x67;
                 message[11] = 0x65;
                 message[12] = 0x73;
                 message[13] = 0x74;
             }
             message_len = 14;
             userID = "1234567812345678";
             userID_len = 16 * 8;
         }
         Crypto_SM2Signature sm2_sig;
         // 清零 r 和 s 数组
         memset(sm2_sig.r, 0, sizeof(sm2_sig.r));  // 将 r 数组初始化为 0
         memset(sm2_sig.s, 0, sizeof(sm2_sig.s));  // 将 s 数组初始化为 0
         printf("\n");
         printf("\nSelect an option:\n");
         printf("1. Generate Signature\n");
         printf("2. Verify Signature\n");
         printf("3. Exit\n");
         printf("Enter your choice (1, 2, or 3): ");
         scanf_s("%d", &mode);
         if (mode == 1) {
             if (Crypto_SM2Sign(&sm2_key, message, message_len, userID, userID_len / 8, &sm2_sig))
             {
                 printf("Signature result:\n");
                 printSignature(&sm2_sig);
             }
             else {
                 printf("Signing failed!\n");
             }
         }
         else {
             printf("Input your signature which you want to verify:\n");
             uint8_t test_signature[2*THIRTY_TRHREE] = { 0 };
             for (int i = 0; i < 2 * THIRTY_TRHREE; i++) {
                 unsigned int byte;
                 if (scanf_s("%2x", &byte) != 1) {
                     printf("Invalid input for signature.\n");
                     return -1;
                 }
                 test_signature[i] = (uint8_t)byte;
             }
             memcpy(sm2_sig.r, test_signature, sizeof(sm2_sig.r));
             memcpy(sm2_sig.s, test_signature+32, sizeof(sm2_sig.s));
             printf("\n");
             if (Crypto_SM2Verify(&sm2_key, message, message_len, userID, userID_len / 8, &sm2_sig))
             {
                 printf("Verification successful!\n");
             }
             else {
                 printf("Validation failed!\n");
             }
         }


         // 释放动态分配的内存
         free(message);
         free(userID);
         printf("\n--------------------------------------------------------------------------------------------------------\n");
     }

     system("pause");
     return 0;
}